package com.leetcode.algorithm.y22.m07;

import java.util.HashSet;
import java.util.LinkedList;
import java.util.Queue;
import java.util.Set;

/**
 * 542. 01 矩阵
 * 
 * https://leetcode.cn/problems/01-matrix/
 * 
 * @author jie.deng
 *
 */
class Question0542Solution01 {
	
	public int[][] updateMatrix(int[][] mat) {
		int rows = mat.length;
		int cols = mat[0].length;
		int[][] grid = new int[rows][cols];

		// 所有0进入队列，然后BFS，距离0的距离为1的进入队列，然后距离为2的进入队列，以此类推
		Queue<int[]> queue = new LinkedList<>();
		for (int i = 0; i < rows; i++) {
			for (int j = 0; j < cols; j++) {
				if (mat[i][j] == 0) {
					queue.offer(new int[] { i, j });
				}
			}
		}

		// dis表示当前点位距离0的距离
		int dis = 1;
		Set<Integer> visitedSet = new HashSet<>();
		while (!queue.isEmpty()) {
			int size = queue.size();
			for (int k = 0; k < size; k++) {
				int[] arr = queue.poll();
				int row = arr[0];
				int col = arr[1];
				// 上
				if (row - 1 >= 0) {
					if (mat[row - 1][col] == 1 && visitedSet.add(getIdx(row - 1, col, cols))) {
						queue.offer(new int[] { row - 1, col });
						grid[row - 1][col] = dis;
					}
				}
				// 下
				if (row + 1 < rows) {
					if (mat[row + 1][col] == 1 && visitedSet.add(getIdx(row + 1, col, cols))) {
						queue.offer(new int[] { row + 1, col });
						grid[row + 1][col] = dis;
					}
				}
				// 左
				if (col - 1 >= 0) {
					if (mat[row][col - 1] == 1 && visitedSet.add(getIdx(row, col - 1, cols))) {
						queue.offer(new int[] { row, col - 1 });
						grid[row][col - 1] = dis;
					}
				}
				// 下
				if (col + 1 < cols) {
					if (mat[row][col + 1] == 1 && visitedSet.add(getIdx(row, col + 1, cols))) {
						queue.offer(new int[] { row, col + 1 });
						grid[row][col + 1] = dis;
					}
				}
			}
			dis++;
		}
		return grid;
	}
	
	private int getIdx(int row, int col, int cols) {
		return row * cols + col;
	}	
    
}